#!/bin/bash
export GLOBAL_LOG_LEVEL=3
export TF_CPP_MIN_LOG_LEVEL=3
export SLOG_PRINT_TO_STDOUT=0
execpath=$(dirname $(realpath $0))

IFS=',' read -ra NPUSARR <<<"${VISIBLE_IDS}"
echo "ADD---${NPUSARR[@]}"

export SOC_VERSION=Ascend910
export HCCL_CONNECT_TIMEOUT=600
export PRINT_MODEL=0
export RANK_SIZE=${#NPUSARR[@]}
export DEVICE_NUM=${#NPUSARR[@]}
export JOB_ID=$DLWS_JOB_ID
export RANK_START=0
export RANK_INDEX=0
export POD_NAME=${DLWS_JOB_ID}

logpath=~/work_dirs/distributed/$DLWS_JOB_ID/$DLWS_ROLE_IDX
mkdir -p $logpath
cd $logpath

if [ $DLWS_NUM_WORKER -gt 1 ]; then
  export POD_NAME=$DLWS_JOB_ID-worker-$DLWS_ROLE_IDX
  export SERVER_ID=${DLWS_ROLE_IDX}
  export RANK_START=$((DEVICE_NUM * SERVER_ID))
  export RANK_SIZE=$(($DLWS_NUM_WORKER * 8))
fi

export RANK_ID=$POD_NAME
ulimit -u unlimited
env

export RANK_TABLE_FILE=/home/$DLWS_USER_NAME/.npu/$DLWS_JOB_ID/hccl_ms.json
rm -rf output kernel_meta
INDEX=0

export RANK_INDEX=$INDEX
for ((device_phy_id = 0; device_phy_id < ${DEVICE_NUM}; device_phy_id++)); do
  export RANK_ID=$((RANK_START + device_phy_id))
  export RANK_INDEX=$((RANK_START + device_phy_id))
  export DEVICE_ID=$device_phy_id
  export DEVICE_INDEX=$device_phy_id
  export RANK_DEVICE_INDEX=$device_phy_id

  python $execpath/train.py --is_distributed=1 --lr=0.01 --data_dir=/data/dataset/storage/coco/2014/raw --pretrained_backbone=/data/dataset/storage/pretrained_models/darknet_weights/darknet53_backbone.ckpt --max_epoch=1 &
  echo "============ CMD ============"
  echo $CMD
  echo "============ CMD ============"
  eval $CMD
done
wait